# Copyright (c) OpenMMLab. All rights reserved.
import random
import warnings

import numpy as np
import torch
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
from mmcv.runner import (HOOKS, DistSamplerSeedHook, EpochBasedRunner,
                         Fp16OptimizerHook, OptimizerHook, build_optimizer,
                         build_runner)
from mmcv.utils import build_from_cfg

from mmdet.core import DistEvalHook, EvalHook
from mmdet.datasets import (build_dataloader, build_dataset,
                            replace_ImageToTensor)
from mmdet.utils import get_root_logger
from easypruner.regularize.sparsity import update_layer , display_layer
@HOOKS.register_module()
class SparsityOptimizerHook(OptimizerHook):
    def __init__(self, prune_ratio, factor, norm_layer_names=None, grad_clip=None):
        super(SparsityOptimizerHook, self).__init__(grad_clip)
        
        self.prune_ratio = prune_ratio
        self.factor = factor
        self.norm_layer_names = norm_layer_names
        self.current_iter = 0

    def after_train_iter(self, runner):

        runner.optimizer.zero_grad()
        runner.outputs['loss'].backward()
        
 	###以下开始maskL1
        ratio = self.prune_ratio  	# 需要设置为你剪枝后想保留的通道数目比率（如，剪枝率为60%，则此处设置0.6）
        factor = self.factor
        norm_layer_names = self.norm_layer_names
        
        update_layer(runner.model, norm_layer_names = norm_layer_names ,  factor=factor ,mask_dict = ratio)
        if self.current_iter % 100 == 0:
            display_layer(runner.model, self.norm_layer_names)
        self.current_iter += 1
        ####

        if self.grad_clip is not None:
            grad_norm = self.clip_grads(runner.model.parameters())
            if grad_norm is not None:
                # Add grad norm to the logger
                runner.log_buffer.update({'grad_norm': grad_norm},
                                         runner.outputs['num_samples'])
        runner.optimizer.step()






def set_random_seed(seed, deterministic=False):
    """Set random seed.

    Args:
        seed (int): Seed to be used.
        deterministic (bool): Whether to set the deterministic option for
            CUDNN backend, i.e., set `torch.backends.cudnn.deterministic`
            to True and `torch.backends.cudnn.benchmark` to False.
            Default: False.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


def train_detector(model,
                   dataset,
                   cfg,
                   distributed=False,
                   validate=False,
                   timestamp=None,
                   meta=None):
    logger = get_root_logger(log_level=cfg.log_level)

    # prepare data loaders
    dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
    if 'imgs_per_gpu' in cfg.data:
        logger.warning('"imgs_per_gpu" is deprecated in MMDet V2.0. '
                       'Please use "samples_per_gpu" instead')
        if 'samples_per_gpu' in cfg.data:
            logger.warning(
                f'Got "imgs_per_gpu"={cfg.data.imgs_per_gpu} and '
                f'"samples_per_gpu"={cfg.data.samples_per_gpu}, "imgs_per_gpu"'
                f'={cfg.data.imgs_per_gpu} is used in this experiments')
        else:
            logger.warning(
                'Automatically set "samples_per_gpu"="imgs_per_gpu"='
                f'{cfg.data.imgs_per_gpu} in this experiments')
        cfg.data.samples_per_gpu = cfg.data.imgs_per_gpu

    runner_type = 'EpochBasedRunner' if 'runner' not in cfg else cfg.runner[
        'type']
    data_loaders = [
        build_dataloader(
            ds,
            cfg.data.samples_per_gpu,
            cfg.data.workers_per_gpu,
            # `num_gpus` will be ignored if distributed
            num_gpus=len(cfg.gpu_ids),
            dist=distributed,
            seed=cfg.seed,
            runner_type=runner_type) for ds in dataset
    ]
    

    
    '''
    以上为原始mmdetection代码
    '''
    # 对网络进行修改，使其改为单变量输入
    from easypruner import  fastpruner
    from easypruner.fastpruner import  getprunelayer
    from easypruner.utils.rebuild import rebuild

    from functools import partial
    import copy
    from mmdet.core.export import pytorch2onnx
    
    normalize_cfg = {'mean': [127,127,127], 'std': [55,55,55]} #格式需要，随便指定即可
    #input_shape = [1,3,1333, 800] # 格式需要，按实际情况指定即可
    #imgpath = "/home/vision/jnf/mmdetection/data/coco/VisDrone/val/0000287_01601_d_0000767.jpg" # 需指定任意一个图片

    input_shape = [1,3,1000, 600] # 格式需要，按实际情况指定即可
    imgpath = "/home/vision/jnf/mmdetection/data/VOCdevkit/VOC2007/JPEGImages/009899.jpg"    # 需指定任意一个图片
    input_config = {  
        'input_shape': input_shape,
        'input_path': imgpath,
        'normalize_cfg': normalize_cfg
    }   
    one_img, one_meta = pytorch2onnx.preprocess_example_input(input_config)

    model.forward = partial(
        model.forward, img_metas=[[one_meta]], return_loss=False)

    
    #加载模型并剪枝   
    state_dict = torch.load( "/home/vision/jnf/mmdetection/work_dirs/baseline/89.7_frozen1_normevalT_faster_rcnn_r50_fpn_1x_VOC_CAR/_epoch_12.pth"
            ) 
    device = next(model.parameters()).device
    model.load_state_dict(state_dict['state_dict'],strict=False)
    model.cpu()
    #获得剪枝层名称
    norm_layer_names = getprunelayer(model)
    print(norm_layer_names)
    model.to(device)
    
    
    #复原模型
    model.forward = model.forward.func 
    '''
    以下为原始mmdetection代码
    '''
    
    


    # put model on gpus
    if distributed:
        find_unused_parameters = cfg.get('find_unused_parameters', False)
        # Sets the `find_unused_parameters` parameter in
        # torch.nn.parallel.DistributedDataParallel
        model = MMDistributedDataParallel(
            model.cuda(),
            device_ids=[torch.cuda.current_device()],
            broadcast_buffers=False,
            find_unused_parameters=find_unused_parameters)
    else:
        model = MMDataParallel(
            model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids)

    # build runner
    optimizer = build_optimizer(model, cfg.optimizer)

    if 'runner' not in cfg:
        cfg.runner = {
            'type': 'EpochBasedRunner',
            'max_epochs': cfg.total_epochs
        }
        warnings.warn(
            'config is now expected to have a `runner` section, '
            'please set `runner` in your config.', UserWarning)
    else:
        if 'total_epochs' in cfg:
            assert cfg.total_epochs == cfg.runner.max_epochs

    runner = build_runner(
        cfg.runner,
        default_args=dict(
            model=model,
            optimizer=optimizer,
            work_dir=cfg.work_dir,
            logger=logger,
            meta=meta))

    # an ugly workaround to make .log and .log.json filenames the same
    runner.timestamp = timestamp
    # fp16 setting
    fp16_cfg = cfg.get('fp16', None)
    if fp16_cfg is not None:
        optimizer_config = Fp16OptimizerHook(
            **cfg.optimizer_config, **fp16_cfg, distributed=distributed)
    elif distributed and 'type' not in cfg.optimizer_config:
        '''
        以上为原代码
        '''
        #注释掉 原始定义 
        #optimizer_config = OptimizerHook(**cfg.optimizer_config)
        optimizer_config = SparsityOptimizerHook(**cfg.optimizer_config, prune_ratio = 0.45 ,  factor=1e-2   ,norm_layer_names= norm_layer_names)

        '''
        以下为原代码
        '''
    else:
        #optimizer_config = cfg.optimizer_config
        optimizer_config = SparsityOptimizerHook(**cfg.optimizer_config, prune_ratio = 0.45,  factor=1e-2   ,norm_layer_names= norm_layer_names)
    # register hooks
    runner.register_training_hooks(cfg.lr_config, optimizer_config,
                                   cfg.checkpoint_config, cfg.log_config,
                                   cfg.get('momentum_config', None))
    if distributed:
        if isinstance(runner, EpochBasedRunner):
            runner.register_hook(DistSamplerSeedHook())

    # register eval hooks
    if validate:
        # Support batch_size > 1 in validation
        val_samples_per_gpu = cfg.data.val.pop('samples_per_gpu', 1)
        if val_samples_per_gpu > 1:
            # Replace 'ImageToTensor' to 'DefaultFormatBundle'
            cfg.data.val.pipeline = replace_ImageToTensor(
                cfg.data.val.pipeline)
        val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
        val_dataloader = build_dataloader(
            val_dataset,
            samples_per_gpu=val_samples_per_gpu,
            workers_per_gpu=cfg.data.workers_per_gpu,
            dist=distributed,
            shuffle=False)
        eval_cfg = cfg.get('evaluation', {})
        eval_cfg['by_epoch'] = cfg.runner['type'] != 'IterBasedRunner'
        eval_hook = DistEvalHook if distributed else EvalHook
        # In this PR (https://github.com/open-mmlab/mmcv/pull/1193), the
        # priority of IterTimerHook has been modified from 'NORMAL' to 'LOW'.
        runner.register_hook(
            eval_hook(val_dataloader, **eval_cfg), priority='LOW')

    # user-defined hooks
    if cfg.get('custom_hooks', None):
        custom_hooks = cfg.custom_hooks
        assert isinstance(custom_hooks, list), \
            f'custom_hooks expect list type, but got {type(custom_hooks)}'
        for hook_cfg in cfg.custom_hooks:
            assert isinstance(hook_cfg, dict), \
                'Each item in custom_hooks expects dict type, but got ' \
                f'{type(hook_cfg)}'
            hook_cfg = hook_cfg.copy()
            priority = hook_cfg.pop('priority', 'NORMAL')
            hook = build_from_cfg(hook_cfg, HOOKS)
            runner.register_hook(hook, priority=priority)

    if cfg.resume_from:
        runner.resume(cfg.resume_from)
    elif cfg.load_from:
        runner.load_checkpoint(cfg.load_from)
    runner.run(data_loaders, cfg.workflow)
